#!/usr/bin/python3

import os
import re

g_gap = -4

class Sequence:
    def __init__(self, name, seq):
        self.name = name
        self.seq = seq
    
    def __len__(self):
        return len(self.seq)
        
def align_sequences(sequences):
    alignments = []
    for i in range(len(sequences) - 1):
        a = sequences[i]
        for j in range(i + 1, len(sequences)):
            b = sequences[j]
            f, c = cost(a, b)
            al_a, al_b = traceback(f, a, b)
            alignments.append((a, b, c, al_a, al_b))
    return alignments

def cost(a, b):
    global g_gap, COST
    m = [[0 for j in range(len(b)+1)] for i in range(len(a)+1)]
    for i in range(0, len(a)+1):
        m[i][0] = i * g_gap
    for j in range(0, len(b)+1):
        m[0][j] = j * g_gap
        
    for i in range(1, len(a)+1):
        for j in range(1, len(b)+1):
            m[i][j] = max(m[i-1][j-1] + COST[a.seq[i-1]][b.seq[j-1]],
                          m[i-1][j]   + COST[a.seq[i-1]]["*"],
                          m[i]  [j-1] + COST["*"][b.seq[j-1]])
    return m, m[len(a)][len(b)]

def traceback(cost, a, b):
    global g_gap, COST
    alignment_a = ""
    alignment_b = ""
    
    i = len(a)
    j = len(b)
    
    while (i > 0 and j > 0):
        score = cost[i][j]
        score_d = cost[i-1][j-1]
        score_l = cost[i-1][j]
        score_u = cost[i][j-1]
        
        if score == score_d + COST[a.seq[i-1]][b.seq[j-1]]:
            alignment_a = a.seq[i-1] + alignment_a
            alignment_b = b.seq[j-1] + alignment_b
            i -= 1
            j -= 1
        elif score == score_l + g_gap:
            alignment_a = a.seq[i-1] + alignment_a
            alignment_b = "-" + alignment_b
            i -= 1
        elif score == score_u + g_gap:
            alignment_a = "-" + alignment_a
            alignment_b = b.seq[j-1] + alignment_b
            j -= 1
    
    for ic in range(i-1, -1, -1):
        alignment_a = a.seq[ic] + alignment_a
        alignment_b = "-" + alignment_b
        
    for jc in range(j-1, -1, -1):
        alignment_a = "-" + alignment_a
        alignment_b = b.seq[jc] + alignment_b
        
    return alignment_a, alignment_b

def make_cost_matrix():
    cost = {}
    pattern = re.compile("[A-Z]|\\*|-?\\d+")
    top = []
    
    file = os.path.join(os.path.dirname(os.path.abspath(__file__)), "testfiles", "BLOSUM62.txt")
    
    count = 0
    with open(file) as f:
        for line in f:
            if line.startswith("#"):
                continue
            
            chars = re.findall(pattern, line)
            if line.startswith(" "):
                top = chars
            else:
                inner= {}
                for c, t in zip(chars[1:], top):
                    inner[t] = int(c)
                cost[chars[0]] = inner                 
                count += 1
    return cost

def parse(filename):
    sequences = []
    with open(filename) as f:
        name = ""
        seq = ""
        for line in f:
            if line.startswith(">"):
                if not seq == "":
                    sequences.append(Sequence(name, seq))
                    name = ""
                    seq = ""
                name = line[1:].rstrip()
            else:
                seq += line.rstrip()
        sequences.append(Sequence(name, seq))
    return sequences
        
if __name__ == "__main__":
    import sys
    
    global COST
    COST = make_cost_matrix()
    
    assert COST["A"]["Z"] == -1
    assert COST["H"]["M"] == -2
    assert COST["P"]["P"] == 7
    assert COST["*"]["G"] == -4
    
    for arg in sys.argv[1:]:
        sequences = parse(arg)
        alignments = align_sequences(sequences)
        for alignment in alignments:
            out = "%s--%s: %d\n" %(alignment[0].name, alignment[1].name, alignment[2])
            out += "%s\n%s" %(alignment[3], alignment[4])
            print(out)